#!/usr/bin/python
"""
Author:			Jimmy Saw
Date modified:	06-19-2012

Description:    This program calculates tetranucleotide frequencies from a given
                multi-fasta file and reports the z score.

Usage:          dissertation_TetraNTCalculatorImproved.py <multi-fasta file> 
                <tetra list>

Example:        dissertation_TetraNTCalculatorImproved.py test-multi.fasta 
                tetra.list

Note: 

z = (x - mu) / rho

tetra.list file should contain a list of tetranucleotide combinations like this:
AAAA
AAAC
AAAG
AAAT
AACA
AACC
AACG
AACT
AAGA
AAGC
.
.
.
TTCG
TTCT
TTGA
TTGC
TTGG
TTGT
TTTA
TTTC
TTTG
TTTT
And should contain a total of 256 combinations.

Note: currently it prints z score in a tab-delimited format like this:
EM7JFSU01D21YQ	-0.471361238918	-0.471361238918	-0.471361238918	3.91658338519
The idea is to use this script to bin metagenomic reads by tetra-nt freq among
other components such as G+C% and other things. I attempt to use z score because
it is a normalized score instead of a raw score which can change based on length
of the sequence.

"""

import re
import sys
import numpy
from Bio import SeqIO

def zscore(x, u, r):
    z = (x - u) / float(r)
    return z

#genome = SeqIO.parse(sys.argv[1], "fasta") #in generator
genome = SeqIO.index(sys.argv[1], "fasta") #in dictionary

tetrafile = open(sys.argv[2], "rU")
tetras = tetrafile.readlines()

tetradict = {}

#for so in genome: #to use with generator
for i, so in genome.iteritems(): #to use with dictionary iteritems
    for t in tetras:
        x = t.strip()
        #tetradict[x] = so.seq.count(x)
        #this statement below counts overlapping tetranucleotides such as GGGGG
        #see this: http://stackoverflow.com/questions/6844005/how-can-i-find-the-number-of-overlapping-sequences-in-a-string-with-python
        tetradict[x] = len(re.findall(r'(?=(%s))' % re.escape(x), so.seq.tostring()))
        #print "Done with", x, tetradict[x]
    tetralist = tetradict.items()
    tetralist.sort()
    total = 0
    numbers = []

    for i in tetralist:
        numbers.append(i[1])
        total = numpy.sum(numbers)
        average = numpy.average(numbers)
        stdev = numpy.std(numbers)
        toprint = so.id

    for j in tetralist:
        z = zscore(j[1], average, stdev)
        toprint += "\t" + '{0:.8}'.format(str(z))
    print toprint


tetrafile.close()
    
